import pandas as pd
import numpy as np
import pycountry

# --- Mapping dictionaries for country codes ---
country_map_2002_manual = {
    'AT': 'AUT','AU': 'AUS','BE-FLA': 'BEL','BG': 'BGR','BR': 'BRA','CH': 'CHE','CL': 'CHL',
    'CY': 'CYP','CZ': 'CZE','DE-E': 'DEU','DE-W': 'DEU','DK': 'DNK','ES': 'ESP','FI': 'FIN',
    'FR': 'FRA','GB-GBN': 'GBR','GB-NIR': 'GBR','HU': 'HUN','IE': 'IRL','IL': 'ISR','JP': 'JPN',
    'LV': 'LVA','MX': 'MEX','NL': 'NLD','NO': 'NOR','NZ': 'NZL','PH': 'PHL','PL': 'POL','PT': 'PRT',
    'RU': 'RUS','SE': 'SWE','SI': 'SVN','SK': 'SVK','TW': 'TWN','US': 'USA'
}

special_map = {
    'DDR': 'DEU','D-W': 'DEU','GB': 'GBR','NIR': 'GBR','RP': 'PHL',
    'CDN': 'CAN','J': 'JPN','BE-WAL': 'BEL','BE-BRU': 'BEL'
}

# --- Attitude recode mappings ---
map_04_egal_int = {3:5, 0:4, 2:3, 1:2, 4:1}
map_04_rev_int  = {3:1, 0:2, 2:3, 1:4, 4:5}
attitude_recode_04 = {'map_egal': map_04_egal_int, 'map_rev': map_04_rev_int}

map_15_egal_int = {1:5,2:4,3:3,4:2,5:1}
map_15_rev_int  = {1:1,2:2,3:3,4:4,5:5}
attitude_recode_15 = {'map_egal': map_15_egal_int, 'map_rev': map_15_rev_int}

# --- Helper functions ---
def final_clean_numeric_robust(series, missing_codes=None, treat_zero_as_missing=False):
    if series is None:
        return pd.Series(dtype='Float64')
    result = pd.to_numeric(series, errors='coerce').astype('Float64')
    if missing_codes:
        result = result.replace(missing_codes, np.nan)
    if treat_zero_as_missing:
        result = result.mask(result == 0.0)
    return result


def clean_spss_categorical_codes(series, missing_codes_numeric=None):
    if series is None:
        return pd.Series(pd.NA, dtype='Int64')
    if not pd.api.types.is_categorical_dtype(series):
        try:
            series = pd.Categorical(series)
        except:
            return pd.Series(pd.NA, index=series.index, dtype='Int64')
    codes = series.codes.astype('Float64')
    codes = codes.replace(-1, np.nan)
    if missing_codes_numeric:
        codes = codes.replace([float(m) for m in missing_codes_numeric], np.nan)
    return codes.astype('Int64')


def get_alpha3_from_numeric(numeric_code):
    if pd.isna(numeric_code):
        return None
    try:
        iso_num = str(int(numeric_code)).zfill(3)
        country = pycountry.countries.get(numeric=iso_num)
        return country.alpha_3 if country else None
    except:
        return None


def standardize_alpha_code(code):
    if pd.isna(code):
        return None
    code = str(code).strip().upper()
    if code in country_map_2002_manual:
        return country_map_2002_manual[code]
    if code in special_map:
        return special_map[code]
    try:
        if len(code) == 3:
            country = pycountry.countries.get(alpha_3=code)
            return country.alpha_3 if country else None
    except:
        pass
    try:
        if len(code) == 2:
            country = pycountry.countries.get(alpha_2=code)
            return country.alpha_3 if country else None
    except:
        pass
    try:
        country = pycountry.countries.lookup(code)
        return country.alpha_3
    except:
        return None
